'''
File: analyzeSequence.py
Author: Adam Hockenberry
Description: 
'''

#Standard path imports
from optparse import OptionParser

#Non-standard imports
from sequenceAnalysisFunctions import *
import glob
import os
import parsingGenbankFiles
from math import exp, log
from matplotlib import pyplot as plt
############################################################################################################
#ignore these two for now, I will try to provide the parameter files
#def create_parameter_dict(directoryOfCodonSumFiles, organismName, dataName):
##For each codon, this will create a dictionary of parameter values from the different fits.
##I suppose it will be an ugly dictionary of dictoinaries
##i.e. first key, codon of interest, second key, which fit, values the parameters for that fit
#    codonDictionary = {}
#    for codonSumFile in glob.glob(directoryToUse + '*codonSums*' + '.dat'):
#        codon = os.path.basename(codonSumFile).split('codonSums')[0]
#        aaSumFile = codonSumFile.replace('codonSums', 'aminoAcidSums')
#        totalLogLDict, totalN, summaryDictionary = fitExpToCDFadam.fitMaxLikelihood(codonSumFile, aaSumFile)
#        codonDictionary[codon] = summaryDictionary
#    if len(codonDictionary.keys()) > 0:
#        writeParametersFile(codonDictionary, organismName, dataName)
#    return codonDictionary
#
#def writeParametersFile(parameterDictionary, organismName, dataName):
#    directory = '../results/' + organismName + '/'
#    if not os.path.exists(directory):
#        os.makedirs(directory)
#    f = open(directory + dataName +'Parameters.txt', 'w')
#    f.write(str(parameterDictionary))
#    f.close()

#############################################################################################################
#read a dictionary that was written according to above, dictionary will be codon:function:parameters
#where exponential = exp and uniform = mean
def read_parameters_file(parameterFileLocation):
    f = open(parameterFileLocation, 'r')
    parameterDictionary = eval(f.read())
    return parameterDictionary


#############################################################################################################
#supporting functions
#############################################################################################################
def geometric_mean(listy, weight=None):
#simple geometric mean
    n = len(listy)
    if weight == None:
        geometric = exp(sum([log(x) for x in listy])/n)
    else:
        listy = [x*y for (x,y) in zip(listy, weight)]
        geometric = exp(sum([log(x) for x in listy])/n)
    return geometric

def calculate_exponential_values(x, parameters):
    a,b,c = parameters
    return a * exp( -1.*b*x) + c 

#############################################################################################################
#main functions
#############################################################################################################


#############################################################################################################
#This is the main function to care about. give me a string and a location for the parameter dictionary to use
#and it will return the CAI valuee, my spatially adjusted version, as well as 2 separate lists which contain
#the weights for each codon for which the CAI is calculated by taking a geometric mean
#############################################################################################################
def main_individual_string(sequenceString, parameterDictionaryLocation):
    parameterDictionary = read_parameters_file(parameterDictionaryLocation)
    CAI, SCAI, weightListCAI, weightListSCAI = get_spatial_CAI(sequenceString, parameterDictionary)
    return CAI, SCAI, weightListCAI, weightListSCAI


#############################################################################################################
#This is a diagnotic plot to use for you to interpret your CAI and SCAI values.
#with the proper inputs it will produce 2 histograms showing CAI and SCAI values for the entire genome
#with green lines denoting the values of the 'highly expressed' reference set that we used
#and a red line showing where your gene falls in comparison, obviously the higher the better
#note: i'm calculating for the genome over and over again, when I'm not so lazy I'll write this to a file
#and for the genome and the gene set it will only need to be calculated once
#############################################################################################################
def main_compare_to_genome(sequenceString, genbankFileLocation, parameterDictionaryLocation, subsetNamesLocation):
    geneCAI, geneSCAI, weightListCAI, weightListSCAI = main_individual_string(sequenceString, parameterDictionaryLocation)
    genomeCAI, genomeSCAI = main_gene_set(genbankFileLocation, parameterDictionaryLocation)
    subsetCAI, subsetSCAI = main_gene_set(genbankFileLocation, parameterDictionaryLocation, subsetNamesLocation)
    plt.figure(1)
    plt.hist(genomeCAI)
    for x in subsetCAI:
        plt.axvline(x, color = 'g', linewidth=1)
    plt.axvline(geneCAI, color = 'r', linewidth=1)
    plt.figure(2)
    plt.hist(genomeSCAI)
    for x in subsetSCAI:
        plt.axvline(x, color = 'g', linewidth=1)
    plt.axvline(geneSCAI, color = 'r', linewidth=1)
    plt.show()
    return

#############################################################################################################
#as mentioned above, this will just iterate through a gene-set and give the CAI/SCAI values
#should just save the values for the genome and common subsets so this only needs to be run once for each
#but haven't yet gotten around to it
#############################################################################################################
def main_gene_set(genbankFileLocation, parameterDictionaryLocation, subset=False):
    genome = parsingGenbankFiles.get_coding_sequences(genbankFileLocation, subset)
    parameterDictionary = read_parameters_file(parameterDictionaryLocation)
    listOfCAIs = []
    listOfSpatialCAIs = []
    for x in range(len(genome)):
        CAI, sCAI, ignore1, ignore2 = get_spatial_CAI(genome[x], parameterDictionary)
        listOfCAIs.append(CAI)
        listOfSpatialCAIs.append(sCAI)
    return listOfCAIs, listOfSpatialCAIs

#############################################################################################################
#the meat of the code, it gets complicated and should be tested more thoroughly for errors
#note that I have a CUTOFF. This will only calculate for codons from 0-500. This can/should be changed
#but right now it is a quirk of the spatial CAI that I can't extend the calculation indefinitely
#############################################################################################################
def get_spatial_CAI(codingSequenceString, parameterDictionary, cutoff=500):
    codingSequenceList = split_string_into_list_of_triplets(codingSequenceString)
    totalWvaluesUniform = []
    totalWvaluesExponential = []
    if codingSequenceList[-1] in ['TAA', 'TGA', 'TAG']:#remove the stop codon
        codingSequenceList = codingSequenceList[:-1]
    for x in range(len(codingSequenceList)):#iterate through each codon apending the weight
        if x < cutoff: ###############How much of the gene to calculate spatial CAI for
            tempUniform, tempExponential = get_W_values(codingSequenceList[x], x, parameterDictionary)
            totalWvaluesUniform.append(tempUniform)
            totalWvaluesExponential.append(tempExponential)
    uniformCAI = geometric_mean(totalWvaluesUniform)
    exponentialCAI = geometric_mean(totalWvaluesExponential)
    return uniformCAI, exponentialCAI, totalWvaluesUniform, totalWvaluesExponential




#############################################################################################################
#The Wvalues are the weights. Each amino acid has a maximum of 1. The RSCU for an individual codon is normalized
#by the maximum for that amino acid to produce the W value
#############################################################################################################
def get_W_values(codon, xValue, parameterDictionary):
    aaDict = get_AA_dict()
####################################################################################################################
#get the conditional probability for the codon of interest at position x, note this doesn't change with x for mean
    probabilityUniform = parameterDictionary[codon]['mean']
    probabilityExponential = calculate_exponential_values(xValue, parameterDictionary[codon]['exp'])#get probability values for the codon
    for aminoAcid in aaDict.keys():
        if codon in aaDict[aminoAcid]:#find the amino acid for the codon of interest
####################################################################################################################
#get probabilities for all codons of the amino acid at that particular position
            exponentialProbabilityList = []
            uniformProbabilityList = []
            for x in aaDict[aminoAcid]:
                tempUniform = parameterDictionary[x]['mean']
                tempExponential = calculate_exponential_values(xValue, parameterDictionary[x]['exp'])
                exponentialProbabilityList.append(tempExponential)
                uniformProbabilityList.append(tempUniform)
            sumProbabilitiesExponential = sum(exponentialProbabilityList)
            sumProbabilitiesUniform = sum(uniformProbabilityList)
####################################################################################################################
#convert all of those probability counts into RSCU values (observed probability relative to expectation)
#where RSCU(codon(x)) = probability(x) / ((1/n) sum(all probabilities for the codon))
            exponentialListRSCU = []
            uniformListRSCU = []
            numberOfCodonsForAA = len(exponentialProbabilityList)
            for x in range(len(exponentialProbabilityList)):
                tempExponential = numberOfCodonsForAA * (exponentialProbabilityList[x]/float(sumProbabilitiesExponential))
                tempUniform = numberOfCodonsForAA * (uniformProbabilityList[x]/float(sumProbabilitiesUniform))
                exponentialListRSCU.append(tempExponential)
                uniformListRSCU.append(tempUniform)
####################################################################################################################
#calculate the RSCU of the codon of interest and use this to determine the weight values by comparing against
#max RSCU of the other possible codons
            actualExponentialRSCU = numberOfCodonsForAA * (probabilityExponential / float(sumProbabilitiesExponential))
            actualUniformRSCU = numberOfCodonsForAA * (probabilityUniform/float(sumProbabilitiesUniform))
            weightExponential = actualExponentialRSCU / max(exponentialListRSCU)
            weightUniform = actualUniformRSCU / max(uniformListRSCU)
####################################################################################################################
#lastly, since the CAI takes a geometric mean, weight values of 0 throw up errors. Set everything below a threshold
#to that threshold
    if weightExponential < 0.001:
        weightExponential = 0.001
    if weightUniform < 0.001:
        weightUniform = 0.001
    return weightUniform, weightExponential

if __name__ == '__main__':
    main_compare_to_genome(gene, blah, blah2, blah3)
